package com.code2roc.xhood.model;

import com.code2roc.xhood.util.ObscureUtil;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;

import java.io.File;
import java.io.FileInputStream;
import java.util.*;
import java.util.stream.Collectors;

public class ClassRelationManage {
    private static HashMap<String,String> classPositionMap = new HashMap<>();
    private static HashMap<String,List<String>> classMap = new HashMap<>();

    public static void init(String outPutPath,File root){
        doInit(outPutPath,root);
        for (List<String> value:classMap.values()) {
            String position = UUID.randomUUID().toString();
            for (String className:value) {
                classPositionMap.put(className,position);
            }
        }
    }

    private static void doInit(String outPutPath,File root){
        if (root.isDirectory()) {
            for (File file : root.listFiles()) {
                doInit(outPutPath,file);
            }
        }
        if (root.getName().endsWith(".class")) {
            List<ClassInfo> classList = new ArrayList<>();
            String className = "";
            try{
                ClassReader cr = new ClassReader(new FileInputStream(root));
                ClassWriter cw = new ClassWriter(cr, 0);
                ClassVisitor cv = new ClassVisitor(Opcodes.ASM9, cw) {

                };
                cr.accept(cv, 0);
                className = cr.getClassName().replace("/",".");
            }catch (Exception e){
                e.printStackTrace();
            }
            findParentClass(outPutPath,className,classList);
            ClassInfo rootClassInfo = classList.stream().filter(a->a.isRoot()).findFirst().get();
            if(!classMap.containsKey(rootClassInfo.className)){
                List<String> contentClass = classList.stream().distinct().map(a->a.getClassName()).collect(Collectors.toList());
                classMap.put(rootClassInfo.className,contentClass);
            }else{
                List<String> dicList = classMap.get(rootClassInfo.className);
                List<String> contentClass = classList.stream().distinct().map(a->a.getClassName()).collect(Collectors.toList());
                dicList.addAll(contentClass);
                dicList = dicList.stream().distinct().collect(Collectors.toList());
                classMap.put(rootClassInfo.className,dicList);
            }
        }
    }

    private static void findParentClass(String outPutPath,String className,List<ClassInfo> childClassList){
        try{
            String filePath = outPutPath + File.separator + "classes" + File.separator + className.replace(".", File.separator) + ".class";
            File file = new File(filePath);
            if (file.exists()) {
                ClassReader cr = new ClassReader(new FileInputStream(new File(filePath)));
                ClassWriter cw = new ClassWriter(cr, 0);
                ClassVisitor cv = new ClassVisitor(Opcodes.ASM9, cw) {

                };
                cr.accept(cv, 0);
                String superClassName = cr.getSuperName().replace("/", ".");
                if(superClassName.equals("java.lang.Object")){
                    childClassList.add(new ClassInfo(className,true));
                }else{
                    childClassList.add(new ClassInfo(className,false));
                    findParentClass(outPutPath,superClassName,childClassList);
                }
            }else{
                childClassList.add(new ClassInfo(className,true));
            }
        }catch (Exception e){
            e.printStackTrace();
        }
    }

    public static HashMap<String, String> getClassPositionMap() {
        return classPositionMap;
    }

    private static class ClassInfo{
        private String className;
        private boolean root;

        public ClassInfo(String className,boolean root){
            this.className = className;
            this.root = root;
        }

        public String getClassName() {
            return className;
        }

        public void setClassName(String className) {
            this.className = className;
        }

        public boolean isRoot() {
            return root;
        }

        public void setRoot(boolean root) {
            this.root = root;
        }
    }
}
